(* $Id: PFOLEncoding.ml 70 2007-10-16 10:05:31Z yann.regisgianas $ *)

(* Pangolin, a functional programming language for correct programs.
   Copyright (C) 2007, Yann Régis-Gianas, François Pottier.
   Contact: yann.regisgianas@gmail.com

   This program is free software: you can redistribute it and/or modify
   it under the terms of the GNU General Public License as published by
   the Free Software Foundation, either version 3 of the License, or
   any later version.

   This program is distributed in the hope that it will be useful,
   but WITHOUT ANY WARRANTY; without even the implied warranty of
   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
   GNU General Public License for more details.

   You should have received a copy of the GNU General Public License
   along with this program.  If not, see <http://www.gnu.org/licenses/>. *)

module X = XCoreSyntax
open XCoreSyntaxExt
open PFOLSyntax
open PFOLSyntaxExt
open PFOLEnv
open PFOLTypesEncoding

exception ErrorWhenEncode of X.lformula

let bind_logic_value pos env x fty = 
  let s = encode_formula_type_scheme_as_object_type_scheme pos env fty in
    bind_value pos env x s

let encode_binding pos (env, locals) (x, ty) = 
  let env, a = bind_logic_value pos env x (formula_monoscheme ty) in
    ((env, Var.AtomSet.add a locals),
     (a, encode_formula_type_as_object_term_type env ty))

let encode_bindings pos (env, locals) bs = 
  Misc.list_fold_map (encode_binding pos) (env, locals) bs

let rec encode_formula_as_predicate pos env locals f = 
  match f.X.fvalue with
    | X.FTrue -> 
	(env, PTrue)

    | X.FFalse ->
	(env, PFalse)

    | X.FVar x -> 
	let y = lookup_predicate env x in
	  if Var.AtomSet.mem y locals then
	    assert false
	  else
	    (env, PApp (y, []))
	      
    | X.FLam abs ->
	assert false

    | X.FForallTys abs ->
	let (vs, t) = X.open_fforalltys_abs abs in
	let tvs = List.map (fun v -> (v, X.KStar)) vs in
	let env = introduce_rigid_type_variables pos env tvs in
	let (env, t) = encode_formula_as_predicate pos env locals t in
	let t = 
	  ForallTys (create_polyf_abs (type_parameters env vs, t)) 
	in
	  (remove_rigid_type_variables env vs, t)

    | X.FForall abs ->
	let (bs, ts, f) = X.open_lforall_abs abs in
	let ((locals', env), bs) = encode_fbindings pos env bs in
	let locals = Var.AtomSet.union locals locals' in
	let (env, t) = encode_formula_as_predicate pos env locals f in
	let ts = List.map (lookup env) ts in
	  (remove_locals locals' env, 
	   Forall (create_forall_abs (bs, ts, t)))

    | X.FExists abs ->
	let (bs, f) = X.open_lexists_abs abs in
	let ((locals', env), bs) = encode_fbindings pos env bs in
	let locals = Var.AtomSet.union locals locals' in
	let (env, t) = encode_formula_as_predicate pos env locals f in
	  (remove_locals locals' env, 
	   Exists (create_exists_abs (bs, t)))

    | X.FEq (ty, f1, f2) ->
	encode_eq pos locals env f1 f2

    | (X.FApp _) as app ->
	let (f, argss) = destruct_fapp f in
	  (match remove_fannot f, argss with
	     | X.FPrimitive X.PNot, [ [ f1 ] ] ->
		 let (env, p1) = 
		   encode_formula_as_predicate pos env locals f1 
		 in
		   (env, Neg p1)

	     | X.FPrimitive ((X.PImply | X.POr | X.PEquiv | X.PAnd) as p),
		 [[ f1; f2 ]] ->
		 let (env, p1) = 
		   encode_formula_as_predicate pos env locals f1 
		 in
		 let (env, p2) = 
		   encode_formula_as_predicate pos env locals f2 
		 in
		   (env, LogicBinOpApp (convert_fprimitive p, p1, p2))

	     | X.FPrimitive
		 ((X.PLessThan | X.PGreaterThan
		  | X.PLessEqualThan | X.PGreaterEqualThan) as p),
		   [[ f1; f2 ]] ->
		 let (env, p1) = encode_formula_as_term pos locals env f1 in
		 let (env, p2) = encode_formula_as_term pos locals env f2 in
		   (env, BinOpApp (convert_fprimitive p, p1, p2))

	     | X.FPrimitive X.Pre, [ [ f1 ]; f2; f3 ] ->
		 let p = match remove_fannot f1.X.fvalue with
		   | X.FVar y -> y
		   | _ -> assert false
		 in
		 let f2 = match f2 with
		   | [ f2 ] -> f2
		   | f2 -> mk_lformula pos (X.FProd f2)
		 in
		 let f3 = match f3 with
		   | [ f3 ] -> f3
		   | f3 -> mk_lformula pos (X.FProd f3)
		 in
		 let (env, a) = encode_formula_as_term pos locals env f2 in
		 let (env, b) = encode_formula_as_term pos locals env f3 in
		   (env, Pre (lookup env p, a, b))

	     | X.FPrimitive X.Post, [ [ f1 ]; f2; f3; f4 ] ->
		 let p = match remove_fannot f1.X.fvalue with
		   | X.FVar y -> y
		   | _ -> assert false
		 in
		 let f2 = match f2 with
		   | [ f2 ] -> f2
		   | f2 -> mk_lformula pos (X.FProd f2)
		 in
		 let f3 = match f3 with
		   | [ f3 ] -> f3
		   | f3 -> mk_lformula pos (X.FProd f3)
		 in
		 let f4 = match f4 with
		   | [ f4 ] -> f4
		   | f4 -> mk_lformula pos (X.FProd f4)
		 in
		 let (env, a) = encode_formula_as_term pos locals env f2 in
		 let (env, b) = encode_formula_as_term pos locals env f3 in
		 let (env, c) = encode_formula_as_term pos locals env f4 in
		   (env, Post (lookup env p, a, b, c))

 	   | X.FVar x, [args] -> 
	       let (env, args) = 
		 encode_formulae_as_terms pos env locals args
	       in
		 (try
		    let y = lookup_predicate env x in
		      if Var.AtomSet.mem y locals then
			assert false
		      else
			(env, PApp (y, args))
		  with Not_found ->
		    (try
		       let y = lookup env x in
			 (env, Reify (PredObjectApp (Var y, args)))
		     with Not_found -> assert false)
		 )
		   
	   | _ -> raise (ErrorWhenEncode (mk_lformula pos app))
	)

    | X.FKApp (f, args) ->
	assert false

    | X.FProd ts ->
	assert false

    | X.FPrimitive p ->
	assert false

    | X.FAnnot (f, ty) ->
	encode_formula_as_predicate pos env locals f

    | _ ->
	assert false

and convert_fprimitive = function
  | X.PAnd -> Conj
  | X.POr -> Disj
  | X.PEquiv -> Equiv
  | X.PImply -> Imply
  | X.PLessThan -> LessThan
  | X.PGreaterThan -> GreaterThan
  | X.PLessEqualThan -> LessEqualThan
  | X.PGreaterEqualThan -> GreaterEqualThan
  | _ -> assert false

and encode_formula_as_term pos locals env f = 
  match f.X.fvalue with
    | X.FPrimitive (X.PEPrimitive (X.PInt x)) ->
	(env, Prim (PInt x))

    | X.FPrimitive (X.PEPrimitive (X.PTrue)) ->
	(env, Prim (PTrueBool))

    | X.FPrimitive (X.PEPrimitive (X.PFalse)) ->
	(env, Prim (PFalseBool))

    | X.FPrimitive (X.PLessThan) ->
	(env, Prim (PLessThanObj))

    | X.FPrimitive (X.PLessEqualThan) ->
	(env, Prim (PLessEqualThanObj))

    | X.FPrimitive (X.PGreaterThan) ->
	(env, Prim (PGreaterThanObj))

    | X.FPrimitive (X.PGreaterEqualThan) ->
	(env, Prim (PGreaterEqualThanObj))

    | X.FVar x -> 
	(try 
	   (env, Var (lookup env x))
	 with Not_found ->
	   (try 
	      (env, Var (lookup_predicate_as_object env x))
	    with Not_found -> 
	      (try 
		 (env, Var (lookup_predicate env x))
	       with Not_found -> 
		 Printf.eprintf "Unbound %s\n" (atom_as_string x);
		 assert false)))

    | X.FProd ts ->
	let (env, ts) = encode_formulae_as_terms pos env locals ts in
	  (env, Prod ts)

    | X.FProd [ t ]  ->
	encode_formula_as_term pos locals env t

    | X.FKApp (k, []) ->
	(env, Var (lookup env k))

    | X.FKApp (k, args) ->
	let (env, ts) = encode_formulae_as_terms pos env locals args in
	  (env, App (lookup env k, ts))
	    
    | X.FApp (f, args) ->
	(match remove_fannot f.X.fvalue, args with
	   | X.FPrimitive ((X.PAdd | X.PSub | X.PMult | X.PSub) as p), 
	     [ t1; t2 ] ->
	       let (env, t1) = encode_formula_as_term pos locals env t1 in
	       let (env, t2) = encode_formula_as_term pos locals env t2 in
		 (env, PrimApp (convert_primitive p, [ t1; t2]))

	   | X.FVar f, args when is_predicate_object env f -> 
	       let (env, ts) = encode_formulae_as_terms pos env locals args in
		 (env, PredObjectApp (Var (lookup env f), ts)) 

	   | X.FVar f, args when not (is_predicate env f) -> 
	       let (env, ts) = encode_formulae_as_terms pos env locals args in
		 (env, App (lookup env f, ts))
		   
	   | _ -> raise (ErrorWhenEncode f))

    | X.FAnnot (f, ty) ->
	let (env, f) = encode_formula_as_term pos locals env f in
	let ty = encode_formula_type_as_term_type env ty in
	  (env, Annot (f, ty))

    | _ -> raise (ErrorWhenEncode f)

and convert_primitive = function
  | X.PAdd -> PAdd
  | X.PMult -> PMult
  | X.PDiv -> PDiv
  | X.PSub -> PSub
  | _ -> assert false

and encode_formulae_as_terms pos env locals fs = 
  Misc.list_fold_map (encode_formula_as_term pos locals) env fs

and encode_fbinding pos (locals, env) (x, ty) = 
  let (env, y) = bind_logic_value pos env x (formula_monoscheme ty) in
  let yty = type_of_id env x in
    ((Var.AtomSet.add y locals, env), (y, yty))

and encode_fbindings pos env bs = 
  Misc.list_fold_map (encode_fbinding pos) (Var.AtomSet.empty, env) bs
    
and remove_locals locals env = 
  Var.AtomSet.fold remove_value locals env

and encode_pre_eq pos locals env c f = 
  let (vs, bs, body) = destruct_poly_flam' c.X.fvalue in
  let args = List.map (fun (x, _) -> X.FVar x) bs in
  let body = mk_fapp pos c.X.fvalue args in
  let body = XCoreNormalizeFormula.normalize body in
    encode_formula_as_predicate pos env locals
      (mk_lformula pos 
	 (mk_fforalltys pos vs
	    (mk_lforall pos bs
	       (mk_fequiv pos 
		  (mk_fapp pos 
		     (mk_fapp pos (X.FPrimitive X.Pre) [ X.FVar f ])
		     [ mk_fprod pos args ])
		  body))))

and encode_post_eq pos locals env c f = 
  let (vs, bs, body)  = destruct_poly_flam' c.X.fvalue in
  let (vs', bs', body) = destruct_poly_flam' body in
  let args  = List.map (fun (x, _) -> X.FVar x) bs in
  let args' = List.map (fun (x, _) -> X.FVar x) bs' in
  let body  = mk_fapp pos (mk_fapp pos c.X.fvalue args) args' in
  let body  = XCoreNormalizeFormula.normalize body in
    encode_formula_as_predicate pos env locals  
      (mk_lformula pos 
	 (mk_fforalltys pos (vs @ vs')
	    ((mk_lforall pos (bs @ bs')
		(mk_fequiv pos 
		   (mk_fapp pos 
		      (mk_fapp pos 
			 (mk_fapp pos (X.FPrimitive X.Post) [X.FVar f])
			 [ mk_fprod pos args ])
		      [ mk_fprod pos args' ])
		   body)))))
    
and encode_eq pos locals env f1 f2 = 
  let ts = destruct_fprod f1.X.fvalue in
  let ts' = destruct_fprod f2.X.fvalue in
    if (Misc.list_same_length ts ts') then
      let args = List.combine ts ts' in
      let (env, eqs) = 
	Misc.list_fold_map 
	  (fun env (t1, t2) -> encode_single_eq pos locals env t1 t2) env args
      in
	(env, mk_conjs eqs)

    else match ts, ts' with
      | [ x ], [ c1; c2 ] 
      | [ c1; c2 ], [ x ] ->
	  let f = as_fvar x.X.fvalue in
(*	  let (env, eq1) = encode_pre_eq pos locals env c1 f in
	  let (env, eq2) = encode_post_eq pos locals env c2 f in *)
	  let (env, eq1) = 
	    encode_eq pos locals env c1 
	      (XCoreNormalizeFormula.lnormalize 
		 ((mk_lformula pos 
		     (mk_fapp pos (X.FPrimitive X.Pre) [x.X.fvalue]) )))
	  in
	  let (env, eq2) = 
	    encode_eq pos locals env c2
	      (XCoreNormalizeFormula.lnormalize 
		 (mk_lformula pos 
		    (mk_fapp pos (X.FPrimitive X.Post) [x.X.fvalue])))
	  in
	    (env, mk_conj eq1 eq2)
      | [x], ts | ts, [x] ->
	  let y = lookup env (as_fvar x.X.fvalue) in
	  let (env, cs) = 
	    Misc.list_fold_map (encode_formula_as_term pos locals) env ts 
	  in
	    (env, PEq (Var y, Prod cs))


(*      | _ -> Printf.eprintf "%d %d\n" 
	  (List.length ts) (List.length ts');
	  assert false *)

and encode_single_eq pos locals env f1 f2 = 
  match (f1.X.fvalue, f2.X.fvalue) with
    | (X.FAnnot (f1, ty1), f2) | (f2, X.FAnnot (f1, ty1)) ->
	(match encode_single_eq pos locals env f1 (mk_lformula pos f2) with
	   | (env, PEq (t1, t2)) ->
	       (* FIXME: assert ty1 = ty2 ? *)
	       let ty = encode_formula_type_as_term_type env ty1 in
		 (env, PEq (Annot (t1, ty), Annot (t2, ty)))
	   | t -> t)

    | X.FLam abs, X.FLam abs' ->
	let ((bs, body), (bs', body')) = X.open2_lfun_abs abs abs' in
	let ((locals', env), bs) = encode_fbindings pos env bs in
	let locals = Var.AtomSet.union locals locals' in
	let (env, eq_t) = encode_single_eq pos locals' env body body' in
	  (remove_locals locals' env, 
	   Forall (create_forall_abs (bs, [], eq_t)))

    | X.FLam abs, t | t, X.FLam abs ->
	let t= mk_lformula pos t in
	let (bs, body) = X.open_lfun_abs abs in
	let args  = List.map (fun (x, _) -> X.FVar x) bs in
	let ((locals', env), bs) = encode_fbindings pos env bs in
	let locals = Var.AtomSet.union locals locals' in
	let t = XCoreNormalizeFormula.lnormalize 
	  (mk_lformula pos (mk_fapp pos t.X.fvalue args)) in
	let (env, eq_t) = encode_single_eq pos locals' env body t in
	  (remove_locals locals' env, 
	   Forall (create_forall_abs (bs, [], eq_t)))

    | _ ->
	try 
	  let (env, t1) = encode_formula_as_term pos locals env f1 in
	  let (env, t2) = encode_formula_as_term pos locals env f2 in
	    (env, PEq (t1, t2))
	with ErrorWhenEncode _ ->
	  let (env, f1) = encode_formula_as_predicate pos env locals f1 in
	  let (env, f2) = encode_formula_as_predicate pos env locals f2 in
	    (env, mk_equiv f1 f2)
	  
let encode_formula_as_predicate pos env f = 
  let f = extrude_foralltys f in
  let (env, f) = encode_formula_as_predicate pos env Var.AtomSet.empty f in
    (env, remove_unused_forall (remove_multiple_annotation f))

